package com.jt1078.streamserver.flv;

import java.util.ArrayList;
import java.util.List;

public class FuaAssembler {
    private byte[] buffer = null;
    private int bufferLen = 0;
    private boolean assembling = false;
    private byte startNalHeader = 0;

    /**
     * 处理单包或FU-A包，返回组装好的NALU（可能为null，需等待更多分包）
     */
    public List<byte[]> feed(byte[] nalu) {
        List<byte[]> out = new ArrayList<>();

        if (nalu.length < 1) return out;
        int nalUnitType = nalu[0] & 0x1F;

        // 单NALU包（非分片）
        if (nalUnitType != 28) {
            // 不是FU-A直接透传
            out.add(nalu);
            clear();
            return out;
        }

        // FU-A分片
        // FU indicator: nalu[0]
        // FU header:    nalu[1]
        // payload:      nalu[2...]
        if (nalu.length < 2) return out;

        byte fuIndicator = nalu[0];
        byte fuHeader = nalu[1];
        boolean start = (fuHeader & 0x80) != 0;
        boolean end = (fuHeader & 0x40) != 0;
        byte nalHeader = (byte)((fuIndicator & 0xE0) | (fuHeader & 0x1F));

        if (start) {
            buffer = new byte[1024 * 128]; // 最多128K一帧，够用
            bufferLen = 0;
            assembling = true;
            startNalHeader = nalHeader;
            // 写起始码+NAL头
            buffer[bufferLen++] = 0x00;
            buffer[bufferLen++] = 0x00;
            buffer[bufferLen++] = 0x00;
            buffer[bufferLen++] = 0x01;
            buffer[bufferLen++] = nalHeader;
            // 写首分包负载
            System.arraycopy(nalu, 2, buffer, bufferLen, nalu.length - 2);
            bufferLen += (nalu.length - 2);
        } else if (assembling) {
            // 追加负载
            System.arraycopy(nalu, 2, buffer, bufferLen, nalu.length - 2);
            bufferLen += (nalu.length - 2);
        }

        if (end && assembling) {
            // 输出组装完整NALU
            byte[] outNalu = new byte[bufferLen];
            System.arraycopy(buffer, 0, outNalu, 0, bufferLen);
            out.add(outNalu);
            clear();
        }
        return out;
    }

    private void clear() {
        buffer = null;
        bufferLen = 0;
        assembling = false;
        startNalHeader = 0;
    }
}
